% coreg_ed.m - find "best" affine transform(s) to fit elp to scalp
%
% using Steven Michael's kdtree routine from:
% http://www.mathworks.com/matlabcentral/fileexchange/
%  7030-kd-tree-nearest-neighbor-and-range-search
%
% 1) build k-d tree for nearest point searc
% 2) find initial closest points
% 3) use "fmincon", "lsqnonlin" or "fminsearch" to find best xform
%
% flags.params; 3=>scale, 6=>shift & rotate, 9=> both
% flags.fit; 1=>fmincon, 2=>lsqnonlin, 3=>fminsearch
%

% function xyz_elp_TR=coreg_ed(xyz_scalp,xyz_elp,flags); 
function xyz_elp_TR=coreg_ed(xyz_scalp);

% no point copying these 10^ times
%
global xyz_elp xyz_elp_TR TreeRoot TreeRoot2 ClosestPts flags
global i_ter trackx0 track_dist

% flags.plot=1;  % 1 => plot, 0 not

% 1) Build the k-d Tree once from the reference datapoints.
%
% disp('Building k-d tree')
if flags.kdtree_type==1
%     TreeRoot=kdtree(xyz_scalp);   % Steven Michael's kdtree routine
    TreeRoot2=KDTreeSearcher(xyz_scalp);
elseif flags.kdtree_type==2;
    % pramod vemulapalli's routine:
    plot_stuff=0; % 0 => don't plot 1 => do
    TreeRoot = kd_buildtree(xyz_scalp,plot_stuff); % pramod vemulapalli's routine
else
    disp('Please set flags.kdtree_type to 1 or 2')
    return
end

% 2) find initial closest points - is this necessary !!! ???
%
% disp('finding Closest Points');
% [Pts_idx,ClosestPts] = kdtree_closestpoint(TreeRoot,xyz_elp);

% 3) use minimization to find best xform
%    N.B. might want to use "lsqnonlin"
%    or other routines from "Minimization" toolbox
%
% disp('Minimizing total distance')

% if flags.params == 3  % 3 params (scale)
%     x0=[1,1,1]; % start with null transform: shift;rot;scale
%     x2=x0';
%     x2=x2(:)';
%     lb=[.8 .8 .8 ];
%     ub=[1.2 1.2 1.2];
% end
% 
% if flags.params == 6  % 6 params (shift,rotate)
%     x0=[0,0,0;0,0,0]; % start with null transform: shift,rotate
%     x2=x0';
%     x2=x2(:)';
%     lb=[-30 -30 -30 -pi/6 -pi/6 -pi/6];
%     ub=[30 30 30 pi/6 pi/6 pi/6];
% end
% 
% if flags.params == 9  % 9 params 
%     x0=[0,0,0;0,0,0;1,1,1]; % start with null transform: shift;rot;scale
%     x2=x0';
%     x2=x2(:)';
%     lb=[-30 -30 -30 -pi/6 -pi/6 -pi/6 .8 .8 .8 ];
%     ub=[30 30 30 pi/6 pi/6 pi/6 1.2 1.2 1.2];
% end
 
i_ter=1;    % initialize iteration counter
tic         % initialize timer

if flags.plot > 1
%     figure
    % plot xyz_scalp
    %
%     scatter3(xyz_scalp(:,1),xyz_scalp(:,2),xyz_scalp(:,3))
%     axis equal
%     axis manual
    hold on

    % plot unaligned sensor points (for testing)
    %
    scatter3(1.02*xyz_elp(:,1),...
             1.02*xyz_elp(:,2),...
             1.02*xyz_elp(:,3),20,'m','filled');
    % axis off
    axis equal
    axis manual
    hold on
    disp('')
    disp('unaligned sensors displayed in magenta:');
    disp('')
    pause(0.1)
end % if flags.plot

if flags.fit == 1 %use fmincon (per Elena)
    disp('Co-registering using fmincon')
    options=optimset('fmincon');% setting the optimization options
    options.LargeScale='off';
    options.Algorithm = 'active-set';

    A= [];b=[];Aeq=[];beq=[]; 
    
    % do rigid body first
    %
    x0=[0,0,0;0,0,0]; % start with null transform: shift,rotate - Elena
%     x0=[0,0,0;0,0,0;0,0,0]; % start with null transform - ed
    x2=x0';
    x2=x2(:)';
    lb=[-30 -30 -30 -pi/6 -pi/6 -pi/6];
    ub=[30 30 30 pi/6 pi/6 pi/6];
    
    [x1,fval, flag, output] = fmincon(@xyz_MRI_fit6,x2,A,b,Aeq,beq,lb,ub,[],options);

    % then scale it
    %
%     x0=[0,0,0;0,0,0;1,1,1]; % start with null transform
%     x2=x0';
%     x2=x2(:)';
%     lb=[.8 .8 .8 ];
%     ub=[1.2 1.2 1.2];
%     
%     [x1,fval, flag, output] = fmincon(@xyz_MRI_fit3,x2,A,b,Aeq,beq,lb,ub,[],options);

    % both at at once ? (fails...)
    %
%         options.TolX=.1;    %[.1, .1, .1, .1, .1, .1, .1, .1, .1];
%         [x1,fval, flag, output] = fmincon(@xyz_MRI_fit9,x2,A,b,Aeq,beq,lb,ub,[],options);

end

if flags.fit == 2
    options=optimset('lsqnonlin');% setting the optimization options
    options.LevenbergMarquardt='on';
%     options.LargeScale='off';
%     options.Display='iter';
%     options.TolPCG=0.01;
    disp('Co-registering using lsqnonlin (L-M)')
    
    % do rigid body first
    %
    [x1,fval, residual, flag, output] = lsqnonlin(@xyz_MRI_fit6,x2,lb,ub,options);
    
    % then scale it
    %
    [x1,fval, residual, flag, output] = lsqnonlin(@xyz_MRI_fit3,x2,lb,ub,options);

    % both at at once ?
    %
%     [x1,fval, residual, flag, output] = lsqnonlin(@xyz_MRI_fit9,x2,lb,ub,options);

end

if flags.fit == 3 % use "fminsearch"
    options = optimset('fminsearch');    % 'Display','Final','TolX',.01);
%     options.Display='iter';
    disp('Co-registering using fminsearch')
    
    % do rigid body first
    %
    [x1,fval,exitflag,output]=fminsearch(@xyz_MRI_fit6,x2,options);

    % then scale it
    %
    [x1,fval,exitflag,output]=fminsearch(@xyz_MRI_fit3,x2,options);
    
    % both at at once ?
    %
%     [x1,fval,exitflag,output]=fminsearch(@xyz_MRI_fit9,x2,options);
    
end

fit_time=toc;
% fprintf('Co-registration time = %5.2f seconds\n',fit_time)

% Write fit results to screen
%
fid=1;
fprintf(fid,'Residual error = %5.2f mm.\n',fval);
fprintf(fid,'Best fit xform = \n');
fprintf(fid,'%8.4f  %8.4f  %8.4f\n',x1);
fprintf(fid,'Fit time = %5.2f\n\n',fit_time);

% Write fit results to text file
%
% TRtxt_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '.txt'];
% % fprintf(fid,'Writing to %s\n',TRtxt_file);
% fid=fopen(TRtxt_file,'w');
% fprintf(fid,'%s\n',TRtxt_file);
% fprintf(fid,'Residual error = %5.2f mm.\n',fval);
% fprintf(fid,'Best fit xform = \n');
% fprintf(fid,'%8.4f  %8.4f  %8.4f\n',x1);
% fprintf(fid,'Fit time = %5.2f\n',fit_time);
% %output.
% %        iterations: 65
% %         funcCount: 902
% %          stepsize: 0.0081
% %      cgiterations: []
% %     firstorderopt: []
% %         algorithm: 'medium-scale: Levenberg-Marquardt, line-search'
% %           message: [1x78 char]
% % out_cell=struct2cell(output);
% % for i=1:2;fprintf(fid,'%s\n',char(out_cell(i,:)));end
% fclose(fid);
return

% define output file names
%
xform_str=num2str(flags.params);
opt_str=num2str(flags.fit);
TRxyz_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '.mat'];
hist_fig_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '_history'];
elpreg_file=[in_path s_fold 'S' subj file_type '_em' xform_str '_O' opt_str '_elpreg'];

% plot fit history
%
annot_str=['S' subj file_type '_em' xform_str '_O' opt_str];
fig1=figure('Position',[0.05*scrsz(3) 0.025*scrsz(4) 0.9* scrsz(3) 0.9*scrsz(4)],...
    'Visible','on');
annotation('textbox',[.25 .95 .5 .05],'EdgeColor','none',...
           'String',annot_str,'Interpreter','none',...
           'FontSize',12,'HorizontalAlignment','center');
for i_plot=1:4
    subplot(2,2,1)
    plot(track_dist)
    xlim([1 length(track_dist)])
    ylim([0 5])
    title('RESIDUAL ERROR')
    subplot(2,2,2)
    plot(trackx0(:,1:3))
    xlim([1 length(track_dist)])
    if flags.params ==3;title('SCALE');else title('SHIFT');end
    if flags.params >3
        subplot(2,2,3)
        plot(trackx0(:,4:6))
        xlim([1 length(track_dist)])
        title('ROTATION')
        if flags.params >6
            subplot(2,2,4)
            plot(trackx0(:,7:9))
            xlim([1 length(track_dist)])
            title('SCALE')
        end
    end
end
% drawnow
print('-dtiff',hist_fig_file)

% 5) display initial fit
%
if flags.plot == 2
%     disp('Plotting initial registration')
    fig2=figure('Position',[0.05*scrsz(3) 0.025*scrsz(4) 0.9* scrsz(3) 0.9*scrsz(4)],...
        'Visible','on');
    subplot(2,3,1)
    plot_MRI_xyz(xyz_scalp,xyz_elp,ClosestPts,'Fiducial (.tai) Alignment',MR_th)
    view(-90,0);
    subplot(2,3,2)
    plot_MRI_xyz(xyz_scalp,xyz_elp,ClosestPts,xyz_file,MR_th)
    view(-90,90);
    subplot(2,3,3)
    plot_MRI_xyz(xyz_scalp,xyz_elp,ClosestPts,'',MR_th)
    view(0,0);
%     drawnow

% 7) display transformed fit
%
%     disp('Plotting coregistered points')
%     figure(fig1)
    subplot(2,3,4)
    plot_MRI_xyz(xyz_scalp,TRxyz,ClosestPts,'Best Fit',MR_th)
    view(-90,0);
    subplot(2,3,5)
    plot_MRI_xyz(xyz_scalp,TRxyz,ClosestPts,TRxyz_file,MR_th)
    view(-90,90);
    subplot(2,3,6)
    plot_MRI_xyz(xyz_scalp,TRxyz,ClosestPts,'',MR_th)
    view(0,0);
%     drawnow
    print('-dtiff',elpreg_file)
%     disp(['Writing ' elpreg_file])
%     close(fig2)
end

% save the original & transformed elp coords
%
clear xyz
disp(['Saving ' TRxyz_file])
xyz=TRxyz;
save(TRxyz_file,'xyz','fit_time','x1','fval','output','trackx0','-mat');

% 8) write transformed elp+acpc.BRIK & HEAD
%
